# -*- coding: utf-8*-

import numpy as np
import pandas as pd
import tushare as ts
import matplotlib as mpl
import matplotlib.pyplot as plt
from datetime import datetime
from matplotlib.dates import DateFormatter
from matplotlib.dates import WeekdayLocator
from matplotlib.dates import MONDAY
from matplotlib.dates import DayLocator
from matplotlib.dates import date2num
from matplotlib.finance import candlestick_ohlc
from matplotlib.font_manager import FontProperties

def get_stock_data():
	stocks=ts.get_stock_basics().index.tolist()
	for code in stocks:
		print('Starting building %s.csv' % code)
		df=ts.get_hist_data(code)
		df.to_csv('d:/tushare_datas/%s.csv' % code)

def get_kdj_data(data,N=0,M=0):
	if N==0:
		N=9
	if M==0:
		M=2
	low_list=data['low'].rolling(N).min()
	low_list.fillna(value=data['low'].expanding().min(),inplace=True)
	high_list=data['high'].rolling(N).max()
	high_list.fillna(value=data['high'].expanding().max(),inplace=True)
	rsv=(data['close']-low_list)/(high_list-low_list)*100
	data['kdj_k']=rsv.ewm(com=M).mean()
	data['kdj_d']=data['kdj_k'].ewm(com=M).mean()
	data['kdj_j']=3*data['kdj_k']-2*data['kdj_d']
	data.fillna(0,inplace=True)
	return data[['kdj_k','kdj_d','kdj_j']]

# def get_ma_data(data,N=0):
# 	if N==0:
# 		N=5
# 	data['ma%d' % N]=data['close'].rolling(N).mean()
# 	data.fillna(0,inplace=True)
# 	return data['ma%d' % N]

def simple_ma(data,window_short=5,window_middle=25,window_long=60):
	data['ma_short']=data['close'].rolling(window_short).mean()
	data['ma_middle']=data['close'].rolling(window_middle).mean()
	data['ma_long']=data['close'].rolling(window_long).mean()
	data['ma_short'].fillna(0,inplace=True)
	data['ma_middle'].fillna(0,inplace=True)
	data['ma_long'].fillna(0,inplace=True)
	data.ix[(data['ma_short'].shift(1)>data['ma_long'].shift(1))&(data['open']<data['close'].shift(1)*1.097),'position']=1
	data.ix[(data['ma_short'].shift(1)<data['ma_long'].shift(1))&(data['open']>data['close'].shift(1)*0.903),'position']=0
	data['position'].fillna(method='ffill',inplace=True)
	data['position'].fillna(0,inplace=True)
	return data[['ma_short','ma_middle','ma_long','position']]

def get_macd_data(data,short=0,long1=0,mid=0):
	if short==0:
		short=12
	if long1==0:
		long1=26
	if mid==0:
		mid=9
	data['sema']=data['close'].ewm(span=short).mean()
	data['lema']=data['close'].ewm(span=long1).mean()
	data.fillna(0,inplace=True)
	data['data_dif']=data['sema']-data['lema']
	data['data_dea']=data['data_dif'].ewm(span=mid).mean()
	data['data_macd']=2*(data['data_dif']-data['data_dea'])
	data.fillna(0,inplace=True)
	return data[['data_dif','data_dea','data_macd']]

def read_csv(code):
	return pd.DataFrame.from_csv('d:/tushare_datas/%s.csv' % code).sort_index()

def sdf(hist_data,code):
	# date1='2016-12-01'
	# date2='2017-12-01'

	# hist_data=ts.get_hist_data('601001',start=date1,end=date2)
	# 对tushare获取到的数据转换成candlestick_ohlc()方法可读取的格式
	data_list = []
	mdate=[]
	for dates,row in hist_data.iterrows():
	    # 将时间转换为数字
	    date_time = datetime.strptime(dates,'%Y-%m-%d')
	    t = date2num(date_time)
	    mdate.append(t)
	    open,high,close,low = row[:4]
	    datas = (t,open,high,low,close)
	    data_list.append(datas)
	 
	# 创建子图
	fig, ax = plt.subplots(1,1,figsize=(10.0,6.0))
	# fig.subplots_adjust(bottom=0.2)
	# 设置X轴刻度为日期时间
	ax.xaxis_date()
	# chinese=FontProperties('宋体')
	plt.xticks(rotation=45)
	plt.yticks()
	plt.title("股票代码：%s K线图" % code,fontproperties=font)
	plt.xlabel("时间",fontproperties=font)
	plt.ylabel("股价（元）",fontproperties=font)
	candlestick_ohlc(ax,data_list,width=0.6,colorup='red',colordown='green')

	# print(data_list[0])
	plt.plot(mdate,hist_data['ma_short'],label='short')
	plt.plot(mdate,hist_data['ma_middle'],label='middle')
	plt.plot(mdate,hist_data['ma_long'],label='long')

	plt.legend()
	plt.grid()
	plt.show()	

font = FontProperties(fname='c:/simsun.ttf',size=12,weight=1)
# dt=read_csv('300059')
# ma=simple_ma(dt)
# print(ma)
# get_stock_data()

date1='2017-01-01'
date2='2017-12-31'
code='300059'
dt=ts.get_hist_data(code,start=date1,end=date2)
dt.sort_index(inplace=True)
simple_ma(dt)
sdf(dt,code)
# pandas_candlestick_ohlc(dt)

